from typing import Dict, List, Optional, Any
from .runtime_info import RuntimeInfo, TaskStatus
from .scheduler_base import SchedulerBase

class TaskAssigner:
    """Manages task assignment and execution tracking in the distributed system.
    
    This class works with the scheduler to assign tasks to workers and monitors
    their execution progress.
    """
    
    def __init__(self, runtime_info: RuntimeInfo, scheduler: SchedulerBase):
        self.runtime_info = runtime_info
        self.scheduler = scheduler
        self.pending_tasks: List[str] = []
        self.max_retries = 3
        self.retry_counts: Dict[str, int] = {}
    
    def assign_pending_tasks(self) -> List[Dict[str, Any]]:
        """Assign pending tasks to available workers.
        
        Returns:
            List[Dict[str, Any]]: List of task assignments
        """
        # Get task assignments from scheduler
        assignments = self.scheduler.assign_tasks()
        
        # Update runtime information for assigned tasks
        for assignment in assignments:
            task_id = assignment['task_id']
            worker_id = assignment['worker_id']
            
            self.runtime_info.update_task_status(
                task_id=task_id,
                status=TaskStatus.RUNNING,
                worker_id=worker_id
            )
            
            if task_id in self.pending_tasks:
                self.pending_tasks.remove(task_id)
        
        return assignments
    
    def handle_task_completion(self, task_id: str, worker_id: str,
                             output_info: Dict[str, Any]) -> None:
        """Handle successful task completion.
        
        Args:
            task_id: Completed task identifier
            worker_id: Worker that completed the task
            output_info: Information about task output (size, location, etc.)
        """
        # Update task status
        self.runtime_info.update_task_status(
            task_id=task_id,
            status=TaskStatus.COMPLETED
        )
        
        # Update scheduler
        self.scheduler.update_status(worker_id, task_id, 'completed')
        
        # Register output data location
        if 'data_id' in output_info:
            self.runtime_info.register_data_location(
                data_id=output_info['data_id'],
                worker_id=worker_id
            )
        
        # Update resource usage
        if 'memory_usage' in output_info:
            self.runtime_info.update_resource_usage(
                worker_id=worker_id,
                memory_usage=output_info['memory_usage']
            )
    
    def handle_task_failure(self, task_id: str, worker_id: str,
                          error_info: Dict[str, Any]) -> None:
        """Handle task execution failure.
        
        Args:
            task_id: Failed task identifier
            worker_id: Worker where the task failed
            error_info: Information about the failure
        """
        # Update retry count
        if task_id not in self.retry_counts:
            self.retry_counts[task_id] = 0
        self.retry_counts[task_id] += 1
        
        if self.retry_counts[task_id] < self.max_retries:
            # Mark task as pending for retry
            self.runtime_info.update_task_status(
                task_id=task_id,
                status=TaskStatus.PENDING
            )
            self.pending_tasks.append(task_id)
        else:
            # Mark task as failed if max retries exceeded
            self.runtime_info.update_task_status(
                task_id=task_id,
                status=TaskStatus.FAILED
            )
        
        # Update scheduler
        self.scheduler.handle_failure(worker_id, task_id)
    
    def get_task_dependencies(self, task_id: str) -> List[Dict[str, Any]]:
        """Get information about task dependencies.
        
        Args:
            task_id: Task identifier
            
        Returns:
            List[Dict[str, Any]]: List of dependency information
        """
        dependencies = self.runtime_info.get_task_dependencies(task_id)
        dependency_info = []
        
        for dep_id in dependencies:
            # Get locations of dependency data
            data_locations = self.runtime_info.get_data_locations(dep_id)
            
            # Get worker loads to help with data transfer decisions
            worker_loads = {
                worker_id: self.runtime_info.get_worker_load(worker_id)
                for worker_id in data_locations
            }
            
            dependency_info.append({
                'task_id': dep_id,
                'data_locations': data_locations,
                'worker_loads': worker_loads
            })
        
        return dependency_info
    
    def can_free_task_data(self, task_id: str) -> bool:
        """Check if a task's output data can be freed.
        
        Args:
            task_id: Task identifier
            
        Returns:
            bool: True if the task's data can be freed
        """
        # Get completed tasks
        completed_tasks = self.runtime_info.get_completed_tasks()
        
        # Check if task is in the list of freeable nodes
        freeable_nodes = self.scheduler.free_nodes(completed_tasks)
        return task_id in freeable_nodes